#! /usr/bin/python2

# filechecksums.py
# generate, store and check file checksums in filesystems'
# extended attributes

# (C)Copyright 2013-2015 by Klaus-J. Wolf, kj (at) seismic (dot) de
# This file is licensed under the General Public License Version 2.
# http://www.gnu.org/licenses/gpl-2.0.html

import os, os.path, sys, hashlib, threading, stat
import xattr


# TODO: -q (quickcheck) -C (copy) -m (metacopy) -M (metacopy&check)
# TODO: -X (export as XML) -I (import from XML)


hashnames=["md5","sha512","ripemd160","whirlpool"]
extras=["mtime","size"]
domprefix="user.hash."

DEBUG=False
VERBOSITY=0
TOTAL=0
COUNT=0
LOG=None


class TaskClass:
    
    def __init__(self, buf, hashes):
        self.buf = buf
        self.hashes = hashes
    
    def start(self):
        threads=[]
        for h in self.hashes:
            thr=threading.Thread(target=self.f, args=(h,))
            thr.start()
            threads.append(thr)
        for thr in threads:
            thr.join()
        for thr in threads:
            del thr
    
    def f(self, h):
        h.update(self.buf)


def gethashes(fn, fndest=None, blksz=0x1000000):
    
    hashv=[]
    for hn in hashnames:
        hashv.append(hashlib.new(hn))
    
    n=0
    f=open(fn, "rb")
    if fndest:
        fdest=open(fndest, "xb")
    else:
        fdest=None
    while True:
        b=f.read(blksz)
        if not b:
            break
        x = TaskClass(b, hashv)
        x.start()
        #for h in hashv:
        #    h.update(b)
        n+=len(b)
        if fdest:
            fdest.write(b)
    f.close()
    if fdest:
        fdest.close()
    
    res = {}
    for i in range(len(hashnames)):
        hn = hashnames[i]
        vh = hashv[i].hexdigest()
        res[hn] = vh
    return res


def getadditionalinfo(fn, res={}):
    st = os.lstat(fn)
    res['mtime'] = str(st.st_mtime).strip()
    res['size'] = str(st.st_size).strip()
    return res


def domainize(old, dom):
    res = {}
    for k,v in old.items():
        res[dom+k] = v
    return res


def readxattrvalues(fn, domain):
    res = {}
    attrs = xattr.listxattr(fn)
    for attr in attrs:
        if attr.startswith(domain):
            res[attr] = xattr.getxattr(fn, attr)
    return res



def hash_add(fn, mode):
    global COUNT
    COUNT+=1
    res = {}
    
    # step 1
    if mode!='o':
        oldhashes = readxattrvalues(fn, domprefix)
    else:
        oldhashes = {}
    
    # step 2
    if mode=='a':
        needscan = False
        for hn1 in hashnames:
            hn = domprefix+hn1
            if hn not in oldhashes:
                needscan = True
                break
        if LOG:
            LOG.write("fn = ")
            LOG.write(fn)
            LOG.write("needscan = "+str(needscan)+"\n")
            LOG.write("oldhashes = "+str(oldhashes)+"\n\n")
        if not needscan:
            return res

    # step 3
    if DEBUG:
        currenthashes = {}
    else:
        currenthashes = gethashes(fn)
    
    if VERBOSITY>1:
        sys.stderr.write("(%5.1f%%)  " % (1e2*COUNT/TOTAL))
    if DEBUG:
        sys.stderr.write("would be scanning: " + fn + "\n")
    elif VERBOSITY>0:
        sys.stderr.write("scanning: " + fn + "\n")
    if VERBOSITY>0:
        sys.stderr.flush()
    
    currenthashes = getadditionalinfo(fn, currenthashes)
    currenthashes = domainize(currenthashes, domprefix)
    
    # step 4
    for hn1 in hashnames+extras:
        hn = domprefix+hn1
        if hn in currenthashes and hn in oldhashes and \
            currenthashes[hn] != oldhashes[hn]:
                res[hn] = (oldhashes[hn], currenthashes[hn])
    
    if mode=='c':
        return res
    
    if mode=='o':
        for hn1 in hashnames+extras:
            hn = domprefix+hn1
            xattr.setxattr(fn, hn, currenthashes[hn])
    elif mode=='a':
        for hn1 in hashnames+extras:
            hn = domprefix+hn1
            if hn in currenthashes and hn not in oldhashes:
                xattr.setxattr(fn, hn, currenthashes[hn])
    elif mode=='u' or mode=='w':
        for hn1 in hashnames+extras:
            hn = domprefix+hn1
            if (hn in currenthashes and hn in oldhashes and \
                currenthashes[hn] != oldhashes[hn]) \
            or hn not in oldhashes:
                    xattr.setxattr(fn, hn, currenthashes[hn])
    
    return res


if __name__=="__main__":

    if len(sys.argv)<2:
        
        print "usage:", sys.argv[0], "-mode|option", "dirs..."
        print "mode = 'a' (add if n/ex) / 'o' (overwrite) / 'u' (update)"
        print "\t'w' (update and warn) / 'c' (check=warn)"
        print "option = 'd' (debug) / 'v' (verbose, repeatable)"
        print "\t'l' (write log)"
        sys.exit(1)

    n = 1
    endofoptions = False
    mode = 'a'
    c = 0

    while n<len(sys.argv):
        
        if not sys.argv[n]:
            continue
        
        if not endofoptions:
            if sys.argv[n][0]=='-':
                if sys.argv[n]=="--":
                    endofoptions = True
                    n+=1
                    continue
                if len(sys.argv[n])>1:
                    for ch in sys.argv[n][1:]:
                        if ch in 'aouwc':
                            mode = ch
                        elif ch=='d':
                            DEBUG=True
                        elif ch=='v':
                            VERBOSITY+=1
                        elif ch=='l':
                            LOG=open("filechecksums.log","at")
                n+=1
                continue

        # pass 1
        if VERBOSITY>1:
            for r,ds,fs in os.walk(sys.argv[n]):
                for fn in fs:
                    ffn = os.path.join(r, fn)
                    try:
                        st = os.lstat(ffn)
                    except OSError as e:
                        continue
                    if stat.S_ISREG(st.st_mode):
                        TOTAL+=1

        # pass 2
        for r,ds,fs in os.walk(sys.argv[n]):
            c += len(fs)
            for fn in fs:
                ffn = os.path.join(r, fn)
                try:
                    st = os.lstat(ffn)
                except OSError as e:
                    continue
                if stat.S_ISREG(st.st_mode):
                    res = hash_add(ffn, mode)
                    if mode in 'cw':
                        for k,v in res:
                            sys.stdout.write(k+": "+v+"\n")
        
        n+=1

    print >>sys.stderr, c, "files done."
